from pydantic import BaseModel, Field, model_validator
from typing import Dict, List, Any, Optional, Union, Literal


__all__ = [
    "OpenaiFunction",
    "OpenaiFunctionCall",
    "OpenaiFunctionCallOption",
    "OpenaiChatCompletionFunctionParameters",
    "OpenaiChatCompletionFunctionParametersProperty",
    "OpenaiChatCompletionFunctionCallOptionParam",
]


class OpenaiChatCompletionFunctionParametersPropertyItems(BaseModel):
    type: str = Field(
        ...,
        pattern="^(string|number|integer|boolean)$",
        description="The type of the item.",
    )


class OpenaiChatCompletionFunctionParametersProperty(BaseModel):
    # type in ["string", "number", "integer", "boolean", "array"]
    type: str = Field(
        ...,
        pattern="^(string|number|integer|boolean|array)$",
        description="The type of the parameter.",
    )

    # items only used in array
    items: Optional[OpenaiChatCompletionFunctionParametersPropertyItems] = Field(
        None,
        description="The items of the parameter. Which is only allowed when type is 'array'.",
    )

    # description should not more than MAXIMUM_PARAMETER_DESCRIPTION_LENGTH characters
    description: str = Field("", max_length=256, description="The description of the parameter.")

    # optional enum
    enum: Optional[List[str]] = Field(
        None,
        description="The enum list of the parameter. Which is only allowed when type is 'string'.",
    )

    @model_validator(mode="before")
    def custom_validate(cls, data: Any):
        if data.get("type") != "string" and data.get("enum"):
            raise ValueError("enum is only allowed when type is 'string'")
        if data.get("type") == "array" and not data.get("items"):
            raise ValueError("items is required when type is 'array'")
        return data


class OpenaiChatCompletionFunctionParameters(BaseModel):
    properties: Dict[str, OpenaiChatCompletionFunctionParametersProperty] = Field(
        ...,
        description="The properties of the parameters.",
    )

    type: str = Field(
        "object",
        Literal="object",
        description="The type of the parameters, which is always 'object'.",
    )

    required: List[str] = Field(
        [],
        description="The required parameters.",
    )

    # check all params in "required" are in properties' keys
    @model_validator(mode="before")
    def validate_required(cls, data: Any):
        if "required" not in data:
            data["required"] = []
        for param in data["required"]:
            if param not in data["properties"]:
                raise ValueError(f"parameter {param} is in required but not in properties")
        return data

    def model_dump(self, **kwargs: Any):
        properties_dict = {
            param_name: param.model_dump(exclude_none=True, **kwargs) for param_name, param in self.properties.items()
        }
        return {
            "type": self.type,
            "properties": properties_dict,
            "required": self.required,
        }


class OpenaiFunction(BaseModel):
    name: str = Field(
        ...,
        description="The name of the function to be called. "
        "Must be a-z, A-Z, 0-9, or contain underscores and dashes, with a maximum length of 64.",
    )
    description: str = Field(
        None,
        description="A description of what the function does, used by the model to choose when and how to call the function.",
    )
    parameters: OpenaiChatCompletionFunctionParameters = Field(
        None,
        description="The parameters the functions accepts, described as a JSON Schema object."
        "See the guide for examples, and the"
        "JSON Schema reference for documentation about the format."
        "Omitting `parameters` defines a function with an empty parameter list.",
    )


class OpenaiChatCompletionFunctionCallOptionParam(BaseModel):
    name: str = Field(..., description="The name of the function to call.")


class OpenaiFunctionCall(BaseModel):
    arguments: str = Field(
        ...,
        description="The arguments to call the function with, as generated by the model in JSON format. Note that the model does not always generate valid JSON, and may hallucinate parameters not defined by your function schema. Validate the arguments in your code before calling your function.",
    )

    name: str = Field(..., description="The name of the function to call.")


OpenaiFunctionCallOption = Union[Literal["none", "auto"], OpenaiChatCompletionFunctionCallOptionParam]
